package cn.seczone.iast.agent;

import cn.seczone.iast.agent.utils.FileUtil;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

/**
 * @author XieShaoping
 */
public class MyTransformer implements ClassFileTransformer {

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        try {
            // 排除一些不需要处理的方法
            //被插桩的类名
            String myClassName = "com/example/webdemo/controller/TestController";
            if (!myClassName.equals(className)) {
                return classfileBuffer;
            }

            ClassReader cr = new ClassReader(classfileBuffer);
            ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
            ClassVisitor cv = new MyClassVisitor(cw);
            cr.accept(cv, ClassReader.EXPAND_FRAMES);
            byte[] bytes = cw.toByteArray();
            FileUtil.save2File("C:\\Users\\xieshaoping\\Desktop\\asm/" + className + ".class", bytes);
            return bytes;
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }

}
